#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
目录存储扫描工具
功能：扫描指定目录，分析存储使用情况，生成包含总体大小、子目录大小饼图和文件类型大小饼图的HTML报告

版本: 1.0.0
作者: 陈振玺
"""

import os
import argparse
import json
from pathlib import Path
from collections import defaultdict
from datetime import datetime
import plotly.express as px
import plotly.graph_objects as go
from plotly.subplots import make_subplots
import plotly.offline as pyo


class DirectoryScanner:
    """
    目录存储扫描器
    
    功能：
    - 扫描指定目录及其子目录
    - 计算总体存储大小
    - 分析子目录存储分布
    - 分析文件类型存储分布
    - 生成交互式HTML报告
    
    版本: 1.0.0
    作者: 陈振玺
    """
    
    def __init__(self, target_directory):
        """
        初始化目录扫描器
        
        Args:
            target_directory (str): 要扫描的目标目录路径
        """
        self.target_directory = Path(target_directory).resolve()
        self.total_size = 0
        self.subdirectory_sizes = defaultdict(int)
        self.file_type_sizes = defaultdict(int)
        self.scan_results = {
            'total_size': 0,
            'subdirectories': {},
            'file_types': {},
            'scan_time': None,
            'target_path': str(self.target_directory)
        }
        
    def format_size(self, size_bytes):
        """
        将字节大小格式化为人类可读的格式
        
        Args:
            size_bytes (int): 字节大小
            
        Returns:
            str: 格式化后的大小字符串
        """
        if size_bytes == 0:
            return "0 B"
        
        size_names = ["B", "KB", "MB", "GB", "TB"]
        i = 0
        while size_bytes >= 1024 and i < len(size_names) - 1:
            size_bytes /= 1024.0
            i += 1
        
        return f"{size_bytes:.2f} {size_names[i]}"
    
    def get_file_extension(self, file_path):
        """
        获取文件扩展名，用于文件类型分类
        
        Args:
            file_path (Path): 文件路径对象
            
        Returns:
            str: 文件扩展名（小写）
        """
        ext = file_path.suffix.lower()
        return ext if ext else '无扩展名'
    
    def scan_directory(self):
        """
        扫描目录，计算各项存储统计信息
        
        Returns:
            dict: 扫描结果字典
        """
        print(f"开始扫描目录: {self.target_directory}")
        
        if not self.target_directory.exists():
            raise FileNotFoundError(f"目录不存在: {self.target_directory}")
        
        if not self.target_directory.is_dir():
            raise NotADirectoryError(f"路径不是目录: {self.target_directory}")
        
        # 重置统计数据
        self.total_size = 0
        self.subdirectory_sizes.clear()
        self.file_type_sizes.clear()
        
        # 遍历目录树
        for root, dirs, files in os.walk(self.target_directory):
            root_path = Path(root)
            
            # 计算当前目录下所有文件的大小
            current_dir_size = 0
            for file_name in files:
                file_path = root_path / file_name
                try:
                    file_size = file_path.stat().st_size
                    current_dir_size += file_size
                    self.total_size += file_size
                    
                    # 统计文件类型大小
                    file_ext = self.get_file_extension(file_path)
                    self.file_type_sizes[file_ext] += file_size
                    
                except (OSError, IOError) as e:
                    print(f"警告: 无法访问文件 {file_path}: {e}")
                    continue
            
            # 统计子目录大小（仅统计直接子目录）
            if root_path.parent == self.target_directory:
                subdir_name = root_path.name
                self.subdirectory_sizes[subdir_name] += current_dir_size
            elif root_path == self.target_directory:
                # 处理根目录下的文件
                if current_dir_size > 0:
                    self.subdirectory_sizes['根目录文件'] += current_dir_size
        
        # 更新扫描结果
        self.scan_results.update({
            'total_size': self.total_size,
            'subdirectories': dict(self.subdirectory_sizes),
            'file_types': dict(self.file_type_sizes),
            'scan_time': datetime.now().isoformat()
        })
        
        print(f"扫描完成，总大小: {self.format_size(self.total_size)}")
        return self.scan_results
    
    def create_subdirectory_pie_chart(self):
        """
        创建子目录大小分布饼图
        
        Returns:
            plotly.graph_objects.Figure: 饼图对象
        """
        if not self.subdirectory_sizes:
            return None
        
        # 准备数据
        labels = list(self.subdirectory_sizes.keys())
        values = list(self.subdirectory_sizes.values())
        
        # 创建饼图
        fig = go.Figure(data=[go.Pie(
            labels=labels,
            values=values,
            textinfo='label+percent+value',
            texttemplate='%{label}<br>%{percent}<br>%{text}',
            text=[self.format_size(v) for v in values],
            hovertemplate='<b>%{label}</b><br>' +
                         '大小: %{text}<br>' +
                         '占比: %{percent}<br>' +
                         '<extra></extra>'
        )])
        
        fig.update_layout(
            title={
                'text': '子目录存储大小分布',
                'x': 0.5,
                'xanchor': 'center',
                'font': {'size': 18}
            },
            font={'size': 12},
            showlegend=True,
            height=500
        )
        
        return fig
    
    def create_file_type_pie_chart(self):
        """
        创建文件类型大小分布饼图
        
        Returns:
            plotly.graph_objects.Figure: 饼图对象
        """
        if not self.file_type_sizes:
            return None
        
        # 准备数据，按大小排序并限制显示数量
        sorted_types = sorted(self.file_type_sizes.items(), key=lambda x: x[1], reverse=True)
        
        # 如果文件类型太多，将小的归类为"其他"
        max_types = 15
        if len(sorted_types) > max_types:
            main_types = sorted_types[:max_types-1]
            other_size = sum(size for _, size in sorted_types[max_types-1:])
            main_types.append(('其他', other_size))
            sorted_types = main_types
        
        labels = [ext if ext else '无扩展名' for ext, _ in sorted_types]
        values = [size for _, size in sorted_types]
        
        # 创建饼图
        fig = go.Figure(data=[go.Pie(
            labels=labels,
            values=values,
            textinfo='label+percent+value',
            texttemplate='%{label}<br>%{percent}<br>%{text}',
            text=[self.format_size(v) for v in values],
            hovertemplate='<b>%{label}</b><br>' +
                         '大小: %{text}<br>' +
                         '占比: %{percent}<br>' +
                         '<extra></extra>'
        )])
        
        fig.update_layout(
            title={
                'text': '文件类型存储大小分布',
                'x': 0.5,
                'xanchor': 'center',
                'font': {'size': 18}
            },
            font={'size': 12},
            showlegend=True,
            height=500
        )
        
        return fig
    
    def generate_html_report(self, output_path):
        """
        生成HTML报告
        
        Args:
            output_path (str): 输出HTML文件路径
        """
        if not self.scan_results['total_size']:
            raise ValueError("请先执行目录扫描")
        
        # 创建子图布局
        fig = make_subplots(
            rows=2, cols=1,
            subplot_titles=('子目录存储大小分布', '文件类型存储大小分布'),
            specs=[[{"type": "pie"}], [{"type": "pie"}]],
            vertical_spacing=0.1
        )
        
        # 添加子目录饼图
        if self.subdirectory_sizes:
            labels1 = list(self.subdirectory_sizes.keys())
            values1 = list(self.subdirectory_sizes.values())
            fig.add_trace(go.Pie(
                labels=labels1,
                values=values1,
                name="子目录",
                textinfo='label+percent',
                hovertemplate='<b>%{label}</b><br>' +
                             '大小: %{text}<br>' +
                             '占比: %{percent}<br>' +
                             '<extra></extra>',
                text=[self.format_size(v) for v in values1]
            ), row=1, col=1)
        
        # 添加文件类型饼图
        if self.file_type_sizes:
            # 处理文件类型数据
            sorted_types = sorted(self.file_type_sizes.items(), key=lambda x: x[1], reverse=True)
            max_types = 15
            if len(sorted_types) > max_types:
                main_types = sorted_types[:max_types-1]
                other_size = sum(size for _, size in sorted_types[max_types-1:])
                main_types.append(('其他', other_size))
                sorted_types = main_types
            
            labels2 = [ext if ext else '无扩展名' for ext, _ in sorted_types]
            values2 = [size for _, size in sorted_types]
            fig.add_trace(go.Pie(
                labels=labels2,
                values=values2,
                name="文件类型",
                textinfo='label+percent',
                hovertemplate='<b>%{label}</b><br>' +
                             '大小: %{text}<br>' +
                             '占比: %{percent}<br>' +
                             '<extra></extra>',
                text=[self.format_size(v) for v in values2]
            ), row=2, col=1)
        
        # 更新布局
        fig.update_layout(
            title={
                'text': f'目录存储分析报告<br><sub>扫描路径: {self.target_directory}</sub>',
                'x': 0.5,
                'xanchor': 'center',
                'font': {'size': 20}
            },
            height=1000,
            showlegend=True,
            font={'size': 12}
        )
        
        # 生成HTML内容
        html_content = f"""
<!DOCTYPE html>
<html lang="zh-CN">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>目录存储分析报告</title>
    <style>
        body {{
            font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
            margin: 0;
            padding: 20px;
            background-color: #f5f5f5;
        }}
        .container {{
            max-width: 1200px;
            margin: 0 auto;
            background-color: white;
            border-radius: 10px;
            box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
            padding: 30px;
        }}
        .header {{
            text-align: center;
            margin-bottom: 30px;
            border-bottom: 2px solid #e0e0e0;
            padding-bottom: 20px;
        }}
        .summary {{
            display: grid;
            grid-template-columns: repeat(auto-fit, minmax(250px, 1fr));
            gap: 20px;
            margin-bottom: 30px;
        }}
        .summary-card {{
            background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
            color: white;
            padding: 20px;
            border-radius: 10px;
            text-align: center;
        }}
        .summary-card h3 {{
            margin: 0 0 10px 0;
            font-size: 1.2em;
        }}
        .summary-card .value {{
            font-size: 2em;
            font-weight: bold;
            margin: 10px 0;
        }}
        .chart-container {{
            margin: 20px 0;
        }}
        .info {{
            background-color: #f8f9fa;
            padding: 15px;
            border-radius: 5px;
            margin-top: 20px;
            border-left: 4px solid #007bff;
        }}
    </style>
</head>
<body>
    <div class="container">
        <div class="header">
            <h1>目录存储分析报告</h1>
            <p><strong>扫描路径:</strong> {self.target_directory}</p>
            <p><strong>扫描时间:</strong> {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}</p>
        </div>
        
        <div class="summary">
            <div class="summary-card">
                <h3>总存储大小</h3>
                <div class="value">{self.format_size(self.total_size)}</div>
            </div>
            <div class="summary-card">
                <h3>子目录数量</h3>
                <div class="value">{len(self.subdirectory_sizes)}</div>
            </div>
            <div class="summary-card">
                <h3>文件类型数量</h3>
                <div class="value">{len(self.file_type_sizes)}</div>
            </div>
        </div>
        
        <div class="chart-container">
            {fig.to_html(include_plotlyjs=True, div_id="charts")}
        </div>
        
        <div class="info">
            <h4>说明</h4>
            <ul>
                <li>总存储大小包含所有子目录和文件的大小</li>
                <li>子目录大小分布显示直接子目录的存储占用情况</li>
                <li>文件类型分布按文件扩展名统计存储占用情况</li>
                <li>如果文件类型超过15种，较小的类型会被归类为"其他"</li>
            </ul>
        </div>
    </div>
</body>
</html>
        """
        
        # 写入HTML文件
        output_path = Path(output_path)
        output_path.parent.mkdir(parents=True, exist_ok=True)
        
        with open(output_path, 'w', encoding='utf-8') as f:
            f.write(html_content)
        
        print(f"HTML报告已生成: {output_path}")
        return str(output_path)


def main():
    """
    命令行主函数
    """
    parser = argparse.ArgumentParser(
        description='目录存储扫描工具 - 分析目录存储使用情况并生成HTML报告',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
示例用法:
  %(prog)s /path/to/directory
  %(prog)s /path/to/directory -o report.html
  %(prog)s /home/user/documents -o /tmp/storage_report.html
        """
    )
    
    parser.add_argument(
        'directory',
        help='要扫描的目录路径'
    )
    
    parser.add_argument(
        '-o', '--output',
        default='directory_storage_report.html',
        help='输出HTML报告文件路径 (默认: directory_storage_report.html)'
    )
    
    parser.add_argument(
        '--version',
        action='version',
        version='目录存储扫描工具 1.0.0 (作者: 陈振玺)'
    )
    
    args = parser.parse_args()
    
    try:
        # 创建扫描器实例
        scanner = DirectoryScanner(args.directory)
        
        # 执行扫描
        results = scanner.scan_directory()
        
        # 生成HTML报告
        output_file = scanner.generate_html_report(args.output)
        
        print(f"\n扫描完成!")
        print(f"总大小: {scanner.format_size(results['total_size'])}")
        print(f"子目录数量: {len(results['subdirectories'])}")
        print(f"文件类型数量: {len(results['file_types'])}")
        print(f"报告文件: {output_file}")
        
    except Exception as e:
        print(f"错误: {e}")
        return 1
    
    return 0


if __name__ == '__main__':
    exit(main())